import os
os.environ['ETS_TOOLKIT'] = 'qt4'#导入chaco时需要
import numpy as np
import scipy.special as sc
from chaco.api import ArrayPlotData, Plot, Legend, cbrewer, palette11, palette14
from chaco.tools.api import DragZoom, PanTool, ZoomTool
from enable.api import ComponentEditor
from traits.api import *
from traitsui.api import *


def full_steam_pressure(T):
    return 100 * np.exp(53.67957 - 6743.769 / T - 4.8451 * np.log(T))


def pressure(h):
    return 101325 * (1 - 2.256e-5 * h) ** 5.256


def temperature(h):
    return 288.15 - 6.5e-3 * h


def N(p, T, f, E):
    return 77.6 / T * (p / 100 + 4810 * f * E / 100 / T)


def refraction_rate(h, f):
    p = pressure(h)
    T = temperature(h)
    E = full_steam_pressure(T)
    return N(p, T, f, E)


# 疑似有误，数值太大
def corrected_refraction_rate(h, f):
    return refraction_rate(h, f) + .157 * h


# Paulus-Jeske
def steam_duct(z, hd=40):
    return corrected_refraction_rate(z, .5) + .125 * (z - hd * np.log(z / 1.5e-4 + 1))

class SteamDuct(HasTraits):
    hd = Range(0, 80,40)
    plot = Instance(Plot)

    def update_plot(self, plot):
        h = np.linspace(0, 300., 301)
        refractionrate=steam_duct(h,self.hd)
        plot.data.update_data(
            h=h,
            refractionrate=refractionrate,
        )

    def _hd_changed(self, event):
        self.update_plot(self.plot)

    def _plot_default(self):
        plot = Plot(ArrayPlotData())
        self.update_plot(plot)
        plot.x_axis.title = "refraction Rate [ Pa ]"
        plot.y_axis.title = "Height [ Pa ]"
        plot.orientation='v'
        plot.plot(("h", "refractionrate"), name='Refraction Rate', color='auto')
        legend = Legend()
        legend.plots = plot.plots
        plot.overlays.append(legend)
        return plot

    traits_view = View(
        Group(
            UCustom(
                "plot",
                editor=ComponentEditor(),
            ),
        ),
        Item('hd', label='波导高度', editor_args={'mode': 'slider'}),
    )


class Refraction_rate(HasTraits):
    temperatureRange = Range(-40, 60)
    plot_refraction_rate = Instance(Plot)

    def update_plot_refraction_rate(self, plot):
        p = np.arange(5000, 101325, 100)
        T = 273.15 + self.temperatureRange
        E = full_steam_pressure(T)
        N0 = N(p, T, 0, E)
        N1 = N(p, T, .2, E)
        N2 = N(p, T, .4, E)
        N3 = N(p, T, .6, E)
        N4 = N(p, T, .8, E)
        N5 = N(p, T, .10, E)
        plot.data.update_data(
            p=p,
            N0=N0,
            N1=N1,
            N2=N2,
            N3=N3,
            N4=N4,
            N5=N5,
        )

    def _temperatureRange_changed(self, event):
        self.update_plot_refraction_rate(self.plot_refraction_rate)

    def _plot_refraction_rate_default(self):
        plot = Plot(ArrayPlotData())
        self.update_plot_refraction_rate(plot)
        plot.title = 'refraction Rate'
        plot.x_axis.title = "Pressure [ Pa ]"
        plot.y_axis.title = "refraction Rate [ Pa ]"
        plot.plot(("p", "N0"), name='Humidity=0', color='auto')
        plot.plot(("p", "N1"), name='Humidity=0.2', color='auto')
        plot.plot(("p", "N2"), name='Humidity=0.4', color='auto')
        plot.plot(("p", "N3"), name='Humidity=0.6', color='auto')
        plot.plot(("p", "N4"), name='Humidity=0.8', color='auto')
        plot.plot(("p", "N5"), name='Humidity=1', color='auto')
        legend = Legend()
        legend.plots = plot.plots
        plot.overlays.append(legend)
        return plot

    traits_view = View(
        Group(
            UCustom(
                "plot_refraction_rate",
                editor=ComponentEditor(),
            ),
        ),
        Item('temperatureRange', label='温度[℃]', editor_args={'mode': 'slider'}),
    )


class BasicRelation(HasTraits):
    plot_test = Instance(Plot)

    def _plot_test_default(self):
        plot = Plot(ArrayPlotData())
        self.clear_plot(plot)
        return plot

    button1 = Button('test1')
    button2 = Button('test2')
    button3 = Button('test3')
    fun = Enum('饱和水汽压', '压强', '温度', '折射率', '修正折射率', '蒸发波导', '函数1', '函数2', '函数3', )

    def _fun_changed(self):
        plot = self.plot_test
        self.clear_plot(plot)
        match self.fun:
            case '函数1':
                x = np.linspace(0, 10, 21)
                y = x
                plot.data.update_data(
                    x=x,
                    y=y,
                    y1=y + 10,
                )
                plot.x_axis.title = "x"
                plot.y_axis.title = "y"
                plot.plot(('x', 'y'), name='fjdk', color='auto')
                plot.plot(('x', 'y1'), name='fjdk1', color='auto')
            case '函数2':
                x = np.linspace(0, 20, 41)
                y = x ** 2
                plot.data.update_data(
                    x=x,
                    y=y,
                    y1=y + 10,
                )
                plot.x_axis.title = "x"
                plot.y_axis.title = "y"
                plot.plot(('x', 'y'), name='fjdk', color='auto')
                plot.plot(('x', 'y1'), name='fjdk1', color='auto')
            case '函数3':
                x = np.linspace(0, 40, 81)
                y = x ** 3
                plot.data.update_data(
                    x=x,
                    y=y,
                    y1=y + 10,
                )
                plot.x_axis.title = "x"
                plot.y_axis.title = "y"
                plot.plot(('x', 'y'), name='fjdk', color='auto')
                plot.plot(('x', 'y1'), name='fjdk1', color='auto')
            case '饱和水汽压':
                T = np.arange(-20, 61)
                plot.data.update_data(
                    temperature=T,
                    full_steam_pressure=full_steam_pressure(T + 273.15),
                )
                plot.title = 'Full Steam Pressure'
                plot.x_axis.title = "Temperature [ K ]"
                plot.y_axis.title = "Full Steam Pressure [ Pa ]"
                plot.plot(("temperature", "full_steam_pressure"), name='Full Steam Pressure', color='auto')
            case '压强':
                height = np.linspace(0, 11000, 111)
                plot.data.update_data(
                    height=height,
                    pressure=pressure(height),
                )
                plot.title = 'Pressure'
                plot.x_axis.title = "Altitide [ m ]"
                plot.y_axis.title = "Pressure [ Pa ]"
                plot.plot(("height", "pressure"), name='Pressure', color='auto')
            case '温度':
                height = np.linspace(0, 11000, 111)
                T = temperature(height)
                plot.data.update_data(
                    height=height,
                    temperature=temperature(height),
                )
                plot.title = 'Pressure'
                plot.x_axis.title = "Altitide [ m ]"
                plot.y_axis.title = "Temperature [ Celcius degree ]"
                plot.plot(("height", "temperature"), name='Temperature', color='auto')
            case '折射率':
                height = np.linspace(0, 11000, 111)
                N0 = refraction_rate(height, 0)
                N1 = refraction_rate(height, .2)
                N2 = refraction_rate(height, .4)
                N3 = refraction_rate(height, .6)
                N4 = refraction_rate(height, .8)
                N5 = refraction_rate(height, 1)
                plot.data.update_data(
                    height=height,
                    N0=N0,
                    N1=N1,
                    N2=N2,
                    N3=N3,
                    N4=N4,
                    N5=N5,
                )
                plot.x_axis.title = "Altitide [ m ]"
                plot.y_axis.title = "refraction Rate"
                plot.plot(("height", "N0"), name='Humidity=0', color='auto')
                plot.plot(("height", "N1"), name='Humidity=0.2', color='auto')
                plot.plot(("height", "N2"), name='Humidity=0.4', color='auto')
                plot.plot(("height", "N3"), name='Humidity=0.6', color='auto')
                plot.plot(("height", "N4"), name='Humidity=0.8', color='auto')
                plot.plot(("height", "N5"), name='Humidity=1', color='auto')
            case '修正折射率':
                height = np.linspace(0, 11000, 111)
                N0 = corrected_refraction_rate(height, 0)
                N1 = corrected_refraction_rate(height, .2)
                N2 = corrected_refraction_rate(height, .4)
                N3 = corrected_refraction_rate(height, .6)
                N4 = corrected_refraction_rate(height, .8)
                N5 = corrected_refraction_rate(height, 1)
                plot.data.update_data(
                    height=height,
                    N0=N0,
                    N1=N1,
                    N2=N2,
                    N3=N3,
                    N4=N4,
                    N5=N5,
                )
                plot.x_axis.title = "Altitide [ m ]"
                plot.y_axis.title = "refraction Rate"
                plot.plot(("height", "N0"), name='Humidity=0', color='auto')
                plot.plot(("height", "N1"), name='Humidity=0.2', color='auto')
                plot.plot(("height", "N2"), name='Humidity=0.4', color='auto')
                plot.plot(("height", "N3"), name='Humidity=0.6', color='auto')
                plot.plot(("height", "N4"), name='Humidity=0.8', color='auto')
                plot.plot(("height", "N5"), name='Humidity=1', color='auto')
            case '蒸发波导':
                height = np.linspace(0, 200, 801)
                M = steam_duct(height)
                plot.data.update_data(
                    height=height,
                    M=M,
                )
                plot.y_axis.title = "Altitide [ m ]"
                plot.x_axis.title = "M"
                plot.orientation = 'v'
                plot.plot(("height", "M"), color='auto')
        legend = Legend()
        legend.plots = plot.plots
        plot.overlays.append(legend)
        plot.request_redraw()

    def traits_init(self):
        self._fun_changed()

    def clear_plot(self, plot):
        plots = [p for p in plot.plots.keys()]
        for p in plots:
            plot.delplot(p)
        for overlay in plot.overlays:
            plot.overlays.remove(overlay)
        plot.request_redraw()

    traits_view = View(
        Group(
            UCustom(
                "plot_test",
                editor=ComponentEditor(),
                resizable=True
            ),
            HGroup(
                UItem('fun'),
                spring,
            ),
        ),
        width=1000,
        height=800,
        resizable=True,
    )


class Form(HasTraits):
    basic_relation = Instance(BasicRelation, ())
    refraction_rate=Instance(Refraction_rate,())
    steamDuct=Instance(SteamDuct,())
    traits_view = View(
        Tabbed(
            UCustom('basic_relation', label='基本参数'),
            UCustom('refraction_rate', label='大气折射率',),
            UCustom('steamDuct', label='蒸发波导',),
        ),
        width=1000,
        height=800,
        resizable=True,
    )


if __name__ == '__main__':
    Form().configure_traits()
